import torch
import triton
import triton.language as tl
from triton.language.math import rsqrt
from typing import Tuple


# grid: (bsz * seq_len, 1, 1)
@triton.jit
def _rms_norm_fwd_kernel(
    X_ptr,
    W_ptr,
    Y_ptr,
    RMS_ptr,
    stride_xr,
    stride_xd,
    stride_wd,
    stride_yr,
    stride_yd,
    stride_rr,
    eps,
    hidden_size: tl.constexpr,
):
    """
    y_i = (x_i / (RMS)) * (offset + wi)
    RMS := sqrt(sum(x_i^2) / N)
    """

    row_idx = tl.program_id(0)

    Y_ptr += row_idx * stride_yr
    X_ptr += row_idx * stride_xr
    RMS_ptr += row_idx * stride_rr

    offs = tl.arange(0, hidden_size)
    x = tl.load(X_ptr + offs * stride_xd)
    x_dtype = x.dtype
    w = tl.load(W_ptr + offs * stride_wd)
    x = x.to(tl.float32)

    rms = rsqrt(tl.sum(x * x, axis=0) / hidden_size + eps)
    tl.store(RMS_ptr, rms)

    x = (x * rms).to(x_dtype)
    y = x * w
    tl.store(Y_ptr + offs * stride_yd, y)


# grid: (sm_cnt, 1, 1)
@triton.jit
def _rms_norm_bwd_kernel(
    X_ptr,
    W_ptr,
    RMS_ptr,
    dX_ptr,
    _dW_ptr,
    dY_ptr,
    stride_xr,
    stride_xd,
    stride_wd,
    stride_rr,
    stride_dxr,
    stride_dxd,
    stride_dwr,
    stride_dyr,
    stride_dyd,
    n_rows,
    hidden_size: tl.constexpr,
    Br: tl.constexpr,
):
    """
    dx = (1 / RMS) * [dy * (w + offset - (1 / N) * (1 / RMS ^ 2) * ((dy * (w + offset)) @ x) * x]
    dw = sum(dy * (x / RMS))
    """

    pid = tl.program_id(0)
    row_start = pid * Br
    row_end = min((pid + 1) * Br, n_rows)

    X_ptr += row_start * stride_xr
    RMS_ptr += row_start * stride_rr
    dX_ptr += row_start * stride_dxr
    _dW_ptr += pid * stride_dwr
    dY_ptr += row_start * stride_dyr

    offs = tl.arange(0, hidden_size)
    dw = tl.zeros((hidden_size,), dtype=tl.float32)
    w = tl.load(W_ptr + offs * stride_wd)

    for i in range(row_start, row_end):
        x = tl.load(X_ptr + offs * stride_xd)
        rms = tl.load(RMS_ptr)
        dy = tl.load(dY_ptr + offs * stride_dyd)
        x_dtype = x.dtype
        x = x.to(tl.float32)

        m = (dy * w).to(tl.float32)
        dx = rms * m
        dx += rms * (-(1 / hidden_size) * rms * rms * tl.sum(m * x, axis=0) * x)
        dw += dy * (x * rms).to(x_dtype)

        tl.store(dX_ptr + offs * stride_dxd, dx.to(x_dtype))

        X_ptr += stride_xr
        RMS_ptr += stride_rr
        dX_ptr += stride_dxr
        dY_ptr += stride_dyr

    tl.store(_dW_ptr + offs * stride_dwr, dw)


def rms_norm_fwd(
    X: torch.Tensor, W: torch.Tensor, eps=1e-6
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """RMS Normalization forward pass

    Args:
        X (torch.Tensor): (bsz, seq_len, hidden_size)
        W (torch.Tensor): (hidden_size, )
        eps: avoid div by 0

    Returns:
        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: (Y, X, RMS)
    """
    bsz, seq_len, hidden_size = X.shape
    assert hidden_size % 2 == 0
    n_rows = bsz * seq_len
    X = X.view(-1, hidden_size)
    Y = torch.empty_like(X)
    RMS = X.new_empty((n_rows,))

    _rms_norm_fwd_kernel[(n_rows,)](
        X_ptr=X,
        W_ptr=W,
        Y_ptr=Y,
        RMS_ptr=RMS,
        stride_xr=X.stride(0),
        stride_xd=X.stride(1),
        stride_wd=W.stride(0),
        stride_yr=Y.stride(0),
        stride_yd=Y.stride(1),
        stride_rr=RMS.stride(0),
        eps=eps,
        hidden_size=hidden_size,
    )
    # save X, RMS for backward
    return Y.view(bsz, seq_len, hidden_size), X, RMS


def rms_norm_bwd(
    dY: torch.Tensor,
    X: torch.Tensor,
    W: torch.Tensor,
    RMS: torch.Tensor,
    inplace: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """RMS Normalization backward pass

    Args:
        dY (torch.Tensor): (bsz, seq_len, hidden_size)
        X (torch.Tensor): (bsz * seq_len, hidden_size)
        W (torch.Tensor): (hidden_size, )
        RMS (torch.Tensor): (bsz * seq_len, )
        inplace (bool): whether to use inplace operation for dX

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: (dX, dW)
    """
    bsz, seq_len, hidden_size = dY.shape
    dY = dY.view(-1, hidden_size)
    n_rows = bsz * seq_len

    sm_cnt = torch.cuda.get_device_properties(X.device).multi_processor_count
    Br = triton.cdiv(n_rows, sm_cnt)
    _dW = W.new_empty((sm_cnt, hidden_size), dtype=torch.float32)
    dX = dY if inplace else torch.empty_like(dY)

    _rms_norm_bwd_kernel[(sm_cnt,)](
        X_ptr=X,
        W_ptr=W,
        RMS_ptr=RMS,
        dX_ptr=dX,
        _dW_ptr=_dW,
        dY_ptr=dY,
        stride_xr=X.stride(0),
        stride_xd=X.stride(1),
        stride_wd=W.stride(0),
        stride_rr=RMS.stride(0),
        stride_dxr=dX.stride(0),
        stride_dxd=dX.stride(1),
        stride_dwr=_dW.stride(0),
        stride_dyr=dY.stride(0),
        stride_dyd=dY.stride(1),
        n_rows=n_rows,
        hidden_size=hidden_size,
        Br=Br,
    )

    dX = dX.view(bsz, seq_len, hidden_size)
    dW = _dW.sum(dim=0).to(W.dtype)
    return dX, dW


class RMSNormTritonFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X: torch.Tensor, W: torch.Tensor, eps=1e-6, inplace=True):
        Y, X, RMS = rms_norm_fwd(X, W, eps)
        ctx.inplace = inplace
        ctx.save_for_backward(X, W, RMS)
        return Y

    @staticmethod
    def backward(ctx, dY: torch.Tensor):
        inplace = ctx.inplace
        X, W, RMS = ctx.saved_tensors
        dX, dW = rms_norm_bwd(dY, X, W, RMS, inplace)
        return dX, dW, None, None


def apply_rms_norm(
    hidden_states: torch.Tensor,
    weight: torch.Tensor,
    eps: float = 1e-6,
    inplace: bool = True,
) -> torch.Tensor:
    return RMSNormTritonFunc.apply(hidden_states, weight, eps, inplace)
